﻿#include "PanelLocator.h"
#include "qtOpenCV/plinesfit_opencv.h"
#include "qtOpenCV/convert.h"
#include <QDebug>

enum DIR {LEFT_TO_RIGHT = 0, RIGHT_TO_LEFT = 1, UP_TO_DOWN = 2, DOWN_TO_UP = 3};

static std::vector<cv::Point2d> edgeExtract(cv::Mat matEdge, DIR dir, cv::Point2d shift)
{
    std::vector<cv::Point2d> edge;

    cv::Mat matTemp;

    switch (dir)
    {
    case LEFT_TO_RIGHT: //LEFT_TO_RIGHT
        for (int j = 0; j < matEdge.rows; j++)
        {
           uchar* current = matEdge.ptr<uchar>(j); // current row
           for (int i = 1; i < matEdge.cols - 1; i++)
           {
               if( current[i] > 128)
               {
                   edge.push_back(cv::Point2d(i + shift.x, j + shift.y));
                   break;
               }
           }
        }
        break;
    case RIGHT_TO_LEFT: //RIGHT_TO_LEFT
        cv::flip(matEdge, matTemp, 0);// 左右翻转一下
        for (int j = 0; j < matTemp.rows; j++)
        {
           uchar* current = matTemp.ptr<uchar>(j); // current row
           for (int i = 1; i < matTemp.cols - 1; i++)
           {
               if( current[i] > 128)
               {
                   edge.push_back(cv::Point2d(matTemp.cols - 1 - i + shift.x, j + shift.y));
                   break;
               }
           }
        }
        break;
    case UP_TO_DOWN: //UP_TO_DOWN
        cv::transpose(matEdge, matTemp);
        for (int j = 0; j < matTemp.rows; j++)
        {
           uchar* current = matTemp.ptr<uchar>(j); // current row
           for (int i = 1; i < matTemp.cols - 1; i++)
           {
               if( current[i] > 128)
               {
                   edge.push_back(cv::Point2d(j + shift.x, i + shift.y));
                   break;
               }
           }
        }
        break;
    case DOWN_TO_UP: //DOWN_TO_UP
        cv::flip(matEdge, matTemp, 1);
        cv::transpose(matTemp, matTemp);
        for (int j = 0; j < matTemp.rows; j++)
        {
           uchar* current = matTemp.ptr<uchar>(j); // current row
           for (int i = 1; i < matTemp.cols - 1; i++)
           {
               if( current[i] > 128)
               {
                   edge.push_back(cv::Point2d(j + shift.x, matTemp.rows - 1 - i + shift.y));
                   break;
               }
           }
        }
        break;

    }
    return edge;
}


PanelLocator::PanelLocator()
{

}

bool PanelLocator::panelDetect(QImage &image)
{
    if(image.isNull()) return false;
    cv::Mat matColorImage = QImage2cvMat(image);
    cv::Mat matGrayImage;
    cv::cvtColor(matColorImage, matGrayImage, cv::COLOR_BGR2GRAY);
    cv::Point pt;

    clear();
    int ret = panelShift(matGrayImage, pt);
    if(ret == 1)
    {
        double theta = 0; //这个变量没什么用。但是 panelRotate 要求有这么一个变量
        panelRotate(matGrayImage, theta);
        return true;
    }
    return false;
}

bool PanelLocator::panelDetect(const cv::Mat &grayImage)
{
    clear();
    cv::Point pt;
    int ret = panelShift(grayImage, pt);
    if(ret == 1)
    {
        double theta = 0; //这个变量没什么用。但是 panelRotate 要求有这么一个变量
        panelRotate(grayImage, theta);
        return true;
    }
    return false;
}

bool PanelLocator::getData(double um_pix, double &x_mm, double &y_mm, double &theta, double &score)
{
    x_mm = 0.0;
    y_mm = 0.0;
    theta = 0.0;
    score = 0.0;
    if(!m_panelLocateFailed)
    {
        x_mm = m_shift.x * um_pix / 1000.0;
        y_mm = m_shift.y * um_pix / 1000.0;
        theta = m_rotateTheta;
        score = m_score;
        return true;
    }
    return false;
}

int PanelLocator::panelShift(const cv::Mat &imageGray, cv::Point2i &s)
{
    cv::Point pos;
    if(imageGray.empty() || m_template.empty() || m_searchRegion.empty())
    {
        return -1;
    }
    m_score = match(imageGray(m_searchRegion), m_template, pos, cv::TM_CCOEFF_NORMED);
    qDebug() << "PanelShift score = " << m_score;
    if(m_score > m_threshold)
    {
        m_panelLocateFailed = false;
        pos += m_searchRegion.tl();
        s = pos - m_panelTopLeft;
        m_shift = s;
        m_shift.x = -m_shift.x;
        return 1;
    }
    else
    {
        m_panelLocateFailed = true;
        s = cv::Point(0, 0);
        m_shift = s;
        return 0;
    }
}

int PanelLocator::panelRotate(const cv::Mat &imageGray, double &theta)
{
    m_rotateTheta = 0;
    theta = 0;
    if(imageGray.empty() || m_template.empty() || m_searchRegion.empty())
    {
        return -1;
    }

    cv::Rect searchRegion = m_lineSearchRegion + m_shift;
    auto edgePoints = edgeExtract(imageGray(searchRegion), LEFT_TO_RIGHT , searchRegion.tl());
    if(edgePoints.size() > 10)
    {
        double a, b, c;
        //lineFit(dg, a, b, c);
        weightedLineFit(edgePoints, a, b, c, 10, 10);
        if(a * b > 0) //直线在 1-3象限，角度为正的
        {
            theta = asin(fabs(b));
            m_rotateTheta = theta;
        }
        else
        {
            theta = -asin(fabs(b));
            m_rotateTheta = theta;
        }

        return 1;
    }
    return 0;
}

void PanelLocator::clear()
{
    m_rotateTheta = 0;
    m_shift = cv::Point2i(0, 0);
    m_score = -1.0;
}

cv::Size PanelLocator::panelSize(const cv::Mat &grayImage)
{
    cv::Mat column, row;
    cv::Mat binImage;
    cv::threshold(grayImage(m_searchRegion), binImage, 0, 255, cv::THRESH_OTSU);
    cv::Mat element = getStructuringElement(0, cv::Size(5, 5)); //构造矩形结构元素
    cv::morphologyEx(binImage, binImage, cv::MORPH_OPEN, element); //去掉小的白色噪点
    cv::reduce(binImage, row, 0, cv::REDUCE_MAX, CV_8U); //reduced to a single row
    cv::reduce(binImage, column, 1, cv::REDUCE_MAX, CV_8U); //reduced to a single column

    int width = cv::countNonZero(row);
    int height = cv::countNonZero(column);

    return cv::Size(width, height);
}

int PanelLocator::panelSizeDeviate(const cv::Mat &grayImage)
{
    cv::Size ps = panelSize(grayImage);
    qDebug() << "size = " << ps.width << ", " << ps.height << "(" << m_panelSize.width << ", " << m_panelSize.height << ")";
    int dev_width = abs(ps.width - m_panelSize.width);
    int dev_height = abs(ps.height - m_panelSize.height);

    return std::max(dev_width, dev_height);
}
